import { OllamaEmbeddings } from "@langchain/community/embeddings/ollama";
import { PDFLoader } from "langchain/document_loaders/fs/pdf";
import { Chroma } from "@langchain/community/vectorstores/chroma";
import { RecursiveCharacterTextSplitter } from "langchain/text_splitter";
import { PromptTemplate } from "@langchain/core/prompts";
import { ChatOllama } from "@langchain/community/chat_models/ollama";
import { createStuffDocumentsChain } from "langchain/chains/combine_documents";
import { StringOutputParser } from "@langchain/core/output_parsers";
import { ChromaClient } from 'chromadb'

const collectionName = "hcy-1"

const clearVectorDB =async ()=>{
    const client = new ChromaClient({
        path: "http://117.72.38.226:8000"
      });
    const colls =await client.listCollections()
    console.log("现有数据集：",colls.length)
    for (let i = 0; i < colls.length; i++) {
        await client.deleteCollection({
            name: colls[i].name
        })
    }
    console.log("删除全部数据集")
    // 创建数据集
    const collection  =await client.getOrCreateCollection({name:collectionName})
    console.log("现有数据集：",collection.length)
}

// 清空数据库，避免干扰数据
clearVectorDB()
// const loader = new PDFLoader("./doc/traingo学习管理系统.pdf");
const loader = new PDFLoader("./doc/愚公移山.pdf");
const docs = await loader.load();
console.log(docs.length)
// 创建一个自定义规则的文档切割器，否则pdf默认就是按自然页切割的
// const textSplitter = new RecursiveCharacterTextSplitter(
//     {
//       chunkSize: 1000,
//       chunkOverlap: 300,
//     }
//     );
// // 对文档进行切割，得到文件块
// const splitDocs = await textSplitter.splitDocuments(docs);
// 实例化embedding模型对象
const embeddings = new OllamaEmbeddings({
    model: "nomic-embed-text", // default value
    baseUrl: "http://117.72.38.226:11434", // default value
    requestOptions: {
    useMMap: true,
    numThread: 6,
    // numGpu: 1,
    },
});
const s1 = Date.now()
// 将文档块和 embedding 传递给 chroma，对数据向量化并存储
await Chroma.fromDocuments(
    docs,
    embeddings,
    {
        collectionName: collectionName, // 
        url: "http://117.72.38.226:8000", // Optional, will default to this value
    }
);
const s2 = Date.now()
console.log("数据向量化耗时：",s2-s1," 毫秒")
// 下面开始rag增强查询，本质就是根据问题在向量数据库中查询对应的文档片段。
// 将片段作为“背景故事”传递给大模型，然后大模型基于这个背景作答

 // 实例化一个用于对话的大模型对象
const qwen = new ChatOllama({
    baseUrl: "http://117.72.38.226:11434", // Default value
    model: "qwen:4b", // Default value
});
 // 给大模型的提示信息，这个提示词的权重很大，
 // 例如，这里如果是让大模型总结，他就会先回答用户问题，然后开始总结，回答的内容很多。 
 // 需求是让模型中介还是
const prompt = PromptTemplate.fromTemplate(
    "请根据文档内容回答用户的问题: {context}"
);
 // 创建一个langchain调用链，将模型和提示词信息对象串联起来
const chain = await createStuffDocumentsChain({
    llm: qwen,
    outputParser: new StringOutputParser(),
    prompt,
});
// 提出问题，这个地方就是读取用户的输入
const question = "金币怎么获取?";
// 实例化一个数据库对象
const vectorStore = await Chroma.fromExistingCollection(
    embeddings, // embedding 模型，要和向量化编码时使用同一模型
    { collectionName: collectionName,
      url: "http://117.72.38.226:8000",
    } // 要查询的数据集，相当于数据库或者表命
);
// 先查询向量库，拿到背景故事,取最符合的5个片段
const storyChunks = await vectorStore.similaritySearch(question,3);
console.log(storyChunks)
const s3 = Date.now()
// 让大模型作答，context中就是本轮回答的背景故事，chain中包含了 对话模型、提示词。
const response=await chain.invoke({
    context: storyChunks,
});
const s4 = Date.now()
console.log(response)
console.log("llm回答问题耗时：",s4-s3," 毫秒")